#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.action import ActionClient
from builtin_interfaces.msg import Duration as DurationMsg
from builtin_interfaces.msg import Time as TimeMsg

from trajectory_msgs.msg import JointTrajectoryPoint
from control_msgs.action import FollowJointTrajectory
from controller_manager_msgs.srv import ListControllers
from rcl_interfaces.srv import GetParameters
from sensor_msgs.msg import JointState


class AutoHome(Node):
    def __init__(self):
        super().__init__('auto_home')

        # 可调参数
        self.declare_parameter('controller_name', 'joint_trajectory_controller')
        self.declare_parameter('joint_names', ['joint11', 'joint2', 'joint3', 'joint4', 'joint5'])   # 需要回零的关节
        self.declare_parameter('target_positions', [0.0, 0.7854, -1.5708, 0.7854, 0.0])         # 对应上面关节的目标值
        self.declare_parameter('duration_sec', 6.0)                    # 回零用时，调大会更慢

        self.controller_name = self.get_parameter('controller_name').get_parameter_value().string_value
        self.home_joint_names = list(self.get_parameter('joint_names').get_parameter_value().string_array_value)
        self.home_target_positions = list(self.get_parameter('target_positions').get_parameter_value().double_array_value)
        self.duration_sec = float(self.get_parameter('duration_sec').value)

        assert len(self.home_joint_names) == len(self.home_target_positions), "joint_names 和 target_positions 长度必须一致"

        # 订阅 joint_states：用来给未回零的关节填“当前值”
        self.js = {}  # joint name -> position
        self.create_subscription(JointState, '/joint_states', self._on_joint_state, 10)

        # 查询控制器状态
        self.cm_cli = self.create_client(ListControllers, '/controller_manager/list_controllers')

        # 读取控制器 joints 参数
        self.param_cli = self.create_client(GetParameters, f'/{self.controller_name}/get_parameters')

        # JTC action 客户端
        self.action_client = ActionClient(self, FollowJointTrajectory,
                                          f'/{self.controller_name}/follow_joint_trajectory')

        self.ctrl_joints = None
        self.sent = False

        # 定时器：循环推进状态机
        self.timer = self.create_timer(0.5, self._tick)

    # 订阅 joint_states
    def _on_joint_state(self, msg: JointState):
        for n, p in zip(msg.name, msg.position):
            self.js[n] = p

    def _tick(self):
        # 1) 等 list_controllers 可用
        if not self.cm_cli.service_is_ready():
            self.get_logger().info('等待 /controller_manager/list_controllers …')
            return

        # 2) 查询控制器是否 active
        req = ListControllers.Request()
        future = self.cm_cli.call_async(req)
        future.add_done_callback(self._on_list_controllers)
        # 停掉当前 tick，等回调继续推进
        self.timer.cancel()

    def _on_list_controllers(self, future):
        try:
            resp = future.result()
        except Exception as e:
            self.get_logger().error(f'查询控制器状态失败: {e}')
            # 过一会儿再试
            self.timer = self.create_timer(0.8, self._tick)
            return

        state = None
        for c in resp.controller:
            if c.name == self.controller_name:
                state = c.state
                break

        if state != 'active':
            self.get_logger().info(f'{self.controller_name} 状态: {state}，继续等待…')
            self.timer = self.create_timer(0.8, self._tick)
            return

        self.get_logger().info(f'{self.controller_name} 已 active，准备读取 joints 参数…')

        # 3) 读取 joints 参数
        if not self.param_cli.service_is_ready():
            self.get_logger().info('等待 get_parameters 服务…')
            self.timer = self.create_timer(0.5, self._tick)
            return

        preq = GetParameters.Request()
        preq.names = ['joints']
        pfut = self.param_cli.call_async(preq)
        pfut.add_done_callback(self._on_got_joints)

    def _on_got_joints(self, future):
        try:
            presp = future.result()
        except Exception as e:
            self.get_logger().error(f'获取 {self.controller_name}.joints 失败: {e}')
            self.timer = self.create_timer(1.0, self._tick)
            return

        if not presp.values or not presp.values[0].string_array_value:
            self.get_logger().error('控制器未返回 joints 列表，检查控制器参数。')
            self.timer = self.create_timer(1.0, self._tick)
            return

        self.ctrl_joints = list(presp.values[0].string_array_value)
        self.get_logger().info(f'控制器关节列表: {self.ctrl_joints}')

        # 4) 等 action server
        if not self.action_client.wait_for_server(timeout_sec=1.0):
            self.get_logger().info('等待 follow_joint_trajectory action server …')
            self.timer = self.create_timer(0.5, self._tick)
            return

        # 5) 发送轨迹
        self._send_home_trajectory()

    def _send_home_trajectory(self):
        if self.sent:
            return

        # 构造 positions：home_joint -> 目标值；其余 -> 当前 joint_states 值（无就用 0）
        home_map = {jn: val for jn, val in zip(self.home_joint_names, self.home_target_positions)}
        positions = []
        for jn in self.ctrl_joints:
            positions.append(home_map.get(jn, self.js.get(jn, 0.0)))

        goal = FollowJointTrajectory.Goal()
        goal.trajectory.joint_names = self.ctrl_joints

        pt = JointTrajectoryPoint()
        pt.positions = positions
        pt.time_from_start = DurationMsg(
            sec=int(self.duration_sec),
            nanosec=int((self.duration_sec % 1.0) * 1e9)
        )
        goal.trajectory.points.append(pt)

        # 给 header.stamp 设置“现在 + 0.2s”
        now = self.get_clock().now().to_msg()
        add_ns = 200_000_000  # 0.2s
        sec = now.sec + (now.nanosec + add_ns) // 1_000_000_000
        nsec = (now.nanosec + add_ns) % 1_000_000_000
        goal.trajectory.header.stamp = TimeMsg(sec=sec, nanosec=nsec)

        # 目标时间容差（可按需调整/也可放到 YAML 的 constraints）
        goal.goal_time_tolerance = DurationMsg(sec=2)

        self.get_logger().info(f'发送回零轨迹：{self.ctrl_joints} -> {positions}，用时 {self.duration_sec:.2f}s')
        send_future = self.action_client.send_goal_async(goal)
        send_future.add_done_callback(self._on_goal_response)
        self.sent = True

    def _on_goal_response(self, future):
        goal_handle = future.result()
        if not goal_handle.accepted:
            self.get_logger().error('回零目标被拒绝（检查 joints 完整性/接口类型/命名是否匹配）。')
            # 允许重试：3 秒后再跑一轮
            self.sent = False
            self.timer = self.create_timer(3.0, self._tick)
            return

        self.get_logger().info('回零目标已接收，等待执行结果…')
        result_future = goal_handle.get_result_async()
        result_future.add_done_callback(self._on_result)

    def _on_result(self, future):
        try:
            result = future.result().result
            self.get_logger().info(f'回零完成，error_code={result.error_code}')
        except Exception as e:
            self.get_logger().error(f'回零执行出错: {e}')
        finally:
            # 一次性任务：完成后退出
            rclpy.shutdown()


def main():
    rclpy.init()
    node = AutoHome()
    rclpy.spin(node)


if __name__ == '__main__':
    main()
